from rdkit import Chem
from rdkit.Chem.Scaffolds import MurckoScaffold

def load_smiles(file_path):
    """
    从文件中加载 SMILES 分子列表96
    :param file_path: 文本文件路径
    :return: SMILES 字符串列表
    """
    smiles_list = []
    with open(file_path, "r") as f:
        for line in f:
            smiles = line.strip()
            if smiles:  # 确保不是空行
                smiles_list.append(smiles)
    return smiles_list


def calculate_scaffold_diversity(smiles_list):
    """
    计算分子的骨架多样性
    :param smiles_list: SMILES 字符串列表
    :return: 骨架数量和骨架多样性比例
    """
    scaffolds = set()
    for smiles in smiles_list:
        mol = Chem.MolFromSmiles(smiles)  # 转换为分子对象
        if mol:  # 确保分子合法
            # 提取骨架
            scaffold = MurckoScaffold.MurckoScaffoldSmiles(mol=mol, includeChirality=False)
            scaffolds.add(scaffold)

    # 计算骨架数量和多样性比例
    unique_scaffold_count = len(scaffolds)
    total_count = len(smiles_list)
    scaffold_diversity_ratio = unique_scaffold_count / total_count

    return unique_scaffold_count, scaffold_diversity_ratio


if __name__ == "__main__":
    # 分子文件路径
    smiles_file = "molecules_3_large_effective.txt"

    # 加载 SMILES 分子
    print("加载 SMILES 分子...")
    smiles_list = load_smiles(smiles_file)

    # 骨架多样性分析
    print("计算骨架多样性...")
    unique_scaffold_count, scaffold_diversity_ratio = calculate_scaffold_diversity(smiles_list)

    # 输出结果
    print(f"总分子数量: {len(smiles_list)}")
    print(f"独特骨架数量: {unique_scaffold_count}")
    print(f"骨架多样性比例: {scaffold_diversity_ratio:.2%}")
